package conn_pool

import (
	"context"
	"errors"
	"fmt"
	"log"
	"net"
	"sync"
	"sync/atomic"
	"time"
)

type idleConn struct {
	conn    net.Conn
	expired time.Time
}

type ReqConn struct {
	conn chan net.Conn
}

type ConnPool struct {
	num          int64                    // 当前连接数
	maxNum       int64                    // 最大连接数
	idleTime     time.Duration            // 空闲时间
	factory      func() (net.Conn, error) // 创建连接方法
	idleConnList chan *idleConn           // 空闲队列 chan的问题：不能从队尾取，队尾过期的概率小
	reqList      chan *ReqConn            // 阻塞的请求队列 不使用chan原因：会阻塞
	mu           sync.Mutex
}

func (c *ConnPool) Get(ctx context.Context) (net.Conn, error) {
	if ctx.Err() != nil {
		return nil, ctx.Err()
	}
	for {
		select {
		case ic, ok := <-c.idleConnList:
			if !ok {
				return nil, fmt.Errorf("队列已关闭")
			}
			if !ic.expired.Before(time.Now()) {
				return ic.conn, nil
			}
			// 连接过期，再取下一个
		//case <-ctx.Done(): // 这段代码不需要，因为有default就不会阻塞;ctx.Done()可以在后续更关键的地方判断
		//	return nil, ctx.Err()
		default:
			// 没有空闲连接，未超期
			c.mu.Lock()
			if c.num >= c.maxNum {
				// 达到最大连接数
				reqConn := &ReqConn{
					conn: make(chan net.Conn, 1),
				}
				c.mu.Unlock()

				select {
				case c.reqList <- reqConn:
					select {
					case conn, ok := <-reqConn.conn:
						if !ok {
							return nil, fmt.Errorf("connPool: 连接已关闭")
						}
						return conn, nil
					case <-ctx.Done():
						close(reqConn.conn) // 从等待队列中移除自己，不然Put方法会放conn进来
						return nil, ctx.Err()
					}
				case <-ctx.Done(): // 因为可能阻塞，所以需要重新判断是否超时
					close(reqConn.conn) // 从等待队列中移除自己，不然Put方法会放conn进来

					return nil, ctx.Err()
				}

			} else {
				conn, err := c.factory()
				if err != nil {
					c.mu.Unlock()
					return nil, err
				}
				c.num++
				c.mu.Unlock()
				return conn, nil
			}

		}
	}
}

func (c *ConnPool) Put(conn net.Conn) error {
	if conn == nil {
		return errors.New("connPool: conn is nil")
	}
	for {
		select {
		case req, ok := <-c.reqList: // 有阻塞的请求，直接唤醒
			if !ok { // 阻塞请求已关闭
				continue
			}
			req.conn <- conn
		default: // 关闭
			select {
			case c.idleConnList <- &idleConn{conn: conn, expired: time.Now().Add(c.idleTime)}: // 放入空闲队列
			default:
				err := conn.Close()
				atomic.AddInt64(&c.num, -1)
				return err
			}

		}
		return nil
	}
}

func (c *ConnPool) Close() {
	close(c.idleConnList)
	close(c.reqList)
}

type Config struct {
	InitNum    int64                    // 初始连接数
	MaxNum     int64                    // 最大连接数
	MaxIdleNum int64                    // 最大空闲连接数
	MaxReqNum  int                      // 最大阻塞请求数
	IdleTime   time.Duration            // 空闲时间
	Factory    func() (net.Conn, error) // 创建连接方法
}

func validate(config *Config) error {
	if config == nil {
		return errors.New("connPool: config is nil")
	}
	ok := config.InitNum <= config.MaxNum && config.InitNum <= config.MaxIdleNum && config.InitNum >= 0
	ok2 := config.MaxIdleNum <= config.MaxNum
	if !(ok && ok2) {
		return errors.New("connPool: 不满足关系 0 <= initNum <= maxIdleNum <= maxNum")
	}
	if config.Factory == nil {
		return errors.New("connPool: factory is nil")
	}
	return nil
}

func NewConnPool(config *Config) (*ConnPool, error) {
	err := validate(config)
	if err != nil {
		return nil, err
	}
	c := &ConnPool{
		maxNum:       config.MaxNum,
		factory:      config.Factory,
		idleConnList: make(chan *idleConn, config.MaxIdleNum),
		idleTime:     config.IdleTime,
		reqList:      make(chan *ReqConn, config.MaxReqNum),
	}

	wg := sync.WaitGroup{}
	wg.Add(int(config.InitNum))
	for i := 0; i < int(config.InitNum); i++ {
		go func() {
			var conn net.Conn
			conn, err = c.factory()
			if err != nil {
				log.Println("connPool: ", err)
				return
			}
			c.idleConnList <- &idleConn{
				conn:    conn,
				expired: time.Now().Add(config.IdleTime),
			}

			atomic.AddInt64(&c.num, 1)
			defer wg.Done()
		}()
	}

	wg.Wait()

	return c, nil
}
